import os
from datetime import datetime

import pandas as pd
from catboost import CatBoostClassifier, Pool
from sklearn.model_selection import StratifiedKFold, cross_val_score
from xbase_util.common_util import date2s

from src.bean.model_database_enum import ModelDatabaseEnum
from src.constant import project_root_path
from src.model.common.sample import predict_sample
from src.model.model_common_util import test_evaluate
from src.util.common_util import while_input, is_int, input_is_yes, printx
from src.util.config_manager import ConfigManager


def common_increasing(config: ConfigManager, model_id, files, df: pd.DataFrame, old_model: dict,
                      cluster: str = "没有聚类"):
    x_resampled, y_resampled, x_test, y_test, test_size, sample_type = predict_sample(df)
    new_model = CatBoostClassifier()
    new_model.load_model(old_model[ModelDatabaseEnum.model_path.value])
    if input_is_yes("[增量训练]是否交叉验证:"):
        n_splits = while_input("[增量训练]请输入交叉验证折数:", is_int)
        printx("[增量训练]使用交叉验证训练模型")
        skf = StratifiedKFold(n_splits=int(n_splits), shuffle=True, random_state=42)
        scores = cross_val_score(new_model, x_resampled, y_resampled, cv=skf, scoring='f1')
        printx(f"[增量训练]交叉验证 F1 分数: {scores}")
        printx(f"[增量训练]平均 F1 分数: {scores.mean():.4f}")
    else:
        printx("[增量训练]未交叉验证")
    printx("[增量训练]使用全部数据训练模型")
    train_pool = Pool(x_resampled, label=y_resampled)
    train_time = date2s(datetime.now())
    new_model.fit(train_pool, init_model=new_model,
                  eval_set=[(x_test, y_test)] if old_model[ModelDatabaseEnum.use_best_model.value] else None)
    accuracy, precision, recall, f1 = test_evaluate(new_model, x_test, y_test)
    if input_is_yes("[增量训练]是否保存模型?"):
        model_name = while_input("[增量训练]请输入模型名称:", None)
        model_description = while_input("[增量训练]请输入模型描述:", None)
        current_model_path = os.path.join(project_root_path, 'model', f"model_{model_id}.pkl")
        new_model.save_model(current_model_path)
        config.database.append({
            ModelDatabaseEnum.id.value: model_id,
            ModelDatabaseEnum.model_path.value: current_model_path,
            ModelDatabaseEnum.model_name.value: model_name,
            ModelDatabaseEnum.model_description.value: model_description,
            ModelDatabaseEnum.train_time.value: train_time,
            ModelDatabaseEnum.model_data_files.value: files + old_model[ModelDatabaseEnum.model_data_files.value],
            ModelDatabaseEnum.increasing_base_on.value: old_model[ModelDatabaseEnum.id.value],
            ModelDatabaseEnum.sample_type.value: '过采样' if f"{sample_type}" == '1' else "欠采样",
            ModelDatabaseEnum.iterations.value: old_model[ModelDatabaseEnum.iterations.value],
            ModelDatabaseEnum.depth.value: old_model[ModelDatabaseEnum.depth.value],
            ModelDatabaseEnum.learning_rate.value: old_model[ModelDatabaseEnum.learning_rate.value],
            ModelDatabaseEnum.test_size.value: test_size,
            ModelDatabaseEnum.f1.value: f1,
            ModelDatabaseEnum.recall.value: recall,
            ModelDatabaseEnum.accuracy.value: accuracy,
            ModelDatabaseEnum.precision.value: precision,
            ModelDatabaseEnum.eval_metric.value: old_model[ModelDatabaseEnum.eval_metric.value],
            ModelDatabaseEnum.use_best_model.value: old_model[ModelDatabaseEnum.use_best_model.value],
            ModelDatabaseEnum.cluster.value: cluster,
        })
        config.db_save()
